#!/usr/bin/env python3
import rclpy
import serial
from typing import Dict, Callable,Optional
import gc
import threading

class Stm32SerialInterface():

    # 帧结构常量
    FRAME_HEADER    =  bytes([0x55, 0xAA])  # 使用bytes替代列表提高比较效率  # 帧头
    FRAME_TAIL      = 0xAA  # 帧尾
    MIN_FRAME_LEN = 6 # 最小帧长检查（帧头(2)+Length(1)+CmdID(1)+CRC(1)+帧尾(1)=6）
    MAX_BUFFER_SIZE  = 1024 # 每次读取大小
    # 命令常量
    CMD_ENCODER     = 0x01  # 速度控制命令
    CMD_BLADE       = 0x02  # 刀盘控制命令
    CMD_STATE       = 0x03  # 状态请求命令

    CMD_STATE_BLADE = 0x70  # 刀盘状态
    CMD_STATE_BATTERY = 0x71 # 电池状态


   
    def __init__(self,node,serial_port:str,serial_baud:int):

        self.serial_lock = threading.Lock()
        
        self.stm32_serial: Optional[serial.Serial] = None
        self.node = node
        self.recv_buffer = bytearray() # 初始化缓冲区

        self.odom_encoder_left = 0  # 初始化编码器值 左
        self.odom_encoder_right = 0 # 初始化编码器值 右

        self.get_uart_data_time = 0.0 # 获取数据时间戳
        self.get_motor_encoder_time = 0.0 # 获取编码器时间戳


        self.command_handlers:Dict[int,Callable] = {
            self.CMD_ENCODER : self._handle_encoder_frame,
            self.CMD_BLADE   : self._handle_blade_frame,
            self.CMD_STATE   : self._handle_state_request_frame
        }

        self.serial_init(serial_port, serial_baud)
        
           
    def serial_init(self,serial_port,serial_baud):
        try:
            self.stm32_serial = serial.Serial(
                port=serial_port,
                baudrate=serial_baud,
                timeout=2,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE
            )
            if self.stm32_serial.is_open:
                self.node.get_logger().info(f"[INFO] Connected to STM32 on {serial_port} at {serial_baud} baud.")
            else:
                self.node.get_logger().error("Port exists but failed to open.")
                raise ConnectionError("Port exists but failed to open.")
        except Exception as e:
            self.node.get_logger().fatal(f"Failed to connect to STM32: {e}")
            raise RuntimeError(f"[ERROR] Failed to connect to STM32: {e}")
    
    def serial_close(self):
        if self.stm32_serial and self.stm32_serial.is_open:
            try:
                self.stm32_serial.close()
                self.node.get_logger().info("Serial port closed")
            except Exception as e:
                self.node.get_logger().error(f"Error closing serial port: {e}")
         
    
    #############
    #上行协议
    #############
    def read_from_serial(self) -> bool:
        """
        从串口读取数据，并将其添加到接收缓冲区中。
        """

        with self.serial_lock:
            # 检查是否已连接到串口
            if not self.stm32_serial or not self.stm32_serial.is_open:
                return False
            
            # 读取可用数据
            avail = self.stm32_serial.in_waiting
            if avail == 0: 
                return False
            # 读取时限制最大长度&防御性读取
            read_size = min(avail, self.MAX_BUFFER_SIZE - len(self.recv_buffer))
            
            try:
                data = self.stm32_serial.read(read_size)
            except serial.SerialException as e:
                self.node.get_logger().error(f"Serial read error: {e}")
                return False
            
            # self.node.get_logger().info(f"read data: {' '.join(f'{b:02X}' for b in data)}")
            self.recv_buffer.extend(data)
            # self.node.get_logger().info(f"recv buffer: {' '.join(f'{b:02X}' for b in self.recv_buffer)}")
            
            
            # 清理溢出数据（仅保留最新的 MAX_BUFFER_SIZE 数据）
            if len(self.recv_buffer) > self.MAX_BUFFER_SIZE:
                del self.recv_buffer[:len(self.recv_buffer)-self.MAX_BUFFER_SIZE]
            
            # 记录获取数据时间戳
            self.get_uart_data_time =  self.node.get_clock().now().nanoseconds / 1e9
        
        # 解析缓冲区        
        return self._parse_recv_buffer()

     
    def _parse_recv_buffer(self) -> bool:
        """
        解析接收缓冲区中的数据，返回是否成功解析了完整的帧。
        """
        processed = False
        start_index = 0  # 跟踪当前处理位置
        recv_len = len(self.recv_buffer)
      
        while recv_len - start_index  >= self.MIN_FRAME_LEN:
            # 在剩余缓冲区中查找帧头
            search_bytes = self.recv_buffer[start_index:]
            header_pos = search_bytes.find(self.FRAME_HEADER)
           
            if header_pos  == -1:
                # 没有找到帧头，清空已扫描部分
                start_index = recv_len  # 清空
                break

            header_pos += start_index  # 转换为绝对位置
            
            # 丢弃帧头前的无效数据
            if header_pos > start_index:
                start_index = header_pos
                continue
           
            if recv_len  - start_index < 3:
                break  # 还没有 Length 字节，等下一轮
          
            # 读取长度字节
            length = self.recv_buffer[start_index + 2] # Length: 从 CmdID 到 CRC 的字节数
            frame_length = 2 + 1 + length + 1  # 总帧长 = 2(帧头) + 1(length) + length + 1(帧尾)
            # self.node.get_logger().info(f"frame_length: {frame_length}")
            if recv_len  - start_index < frame_length:
                # 如果缓冲区中不足以存储完整的帧，则退出循环,等待下一次读取
                break
            
            # 提取完整帧
            frame_start = start_index
            frame_end = frame_start + frame_length
            frame = self.recv_buffer[frame_start:frame_end]
          
       
            # self.node.get_logger().info(f"frame: {' '.join(f'{b:02X}' for b in frame)}")
            
            # frame_header = frame[:2]
            # frame_length = frame[2]
            frame_payload = frame[3:-2] # cmdID data ,不包括帧头、帧尾和CRC 
            frame_crc = frame[-2] # CRC
            frame_tail = frame[-1] # 帧尾
            
            # 校验帧
            if not self._is_valid_frame(frame_payload, frame_crc, frame_tail):
                # self.node.get_logger().info(f"[WARN] Invalid frame: {' '.join(f'{b:02X}' for b in frame)}")
                start_index += 1  # 跳过错误帧头
                continue

            # 分发处理有效帧
            self._handle_frame(frame_payload)  # 传入 CmdID + 数据（不含 CRC）
          
            start_index = frame_end  # 移动到下一帧
       
            processed = True
        
        # 更新接收缓冲区：保留未处理数据
        if start_index > 0:
            self.recv_buffer = self.recv_buffer[start_index:]

        return processed
    
    def _calc_checksum(self, data: bytes) -> int:
        # 简单异或校验
        crc = 0
        for b in data:
            crc ^= b
        return crc
    
    def _is_valid_frame(self, frame_payload:bytes, frame_crc:int, frame_tail:int) -> bool:
        
        # 检查帧尾
        if frame_tail != self.FRAME_TAIL:
            return False
        # 检查 CRC
        calculated_crc = self._calc_checksum(frame_payload)
        received_crc = frame_crc

        # self.node.get_logger().warn(f'calculated_crc: {calculated_crc:#02x},received_crc: {received_crc:#02x}')
  
        return True

        return calculated_crc == received_crc
    
    

    def _handle_frame(self, payload: bytes):
        if not payload:
            return
        
        cmd_id, cmd_data = payload[0], payload[1:]

        handler = self.command_handlers.get(cmd_id)

        if handler:
            handler(cmd_data)
        else:
            self._handle_unknown_command(cmd_id, cmd_data)
        

    def _handle_encoder_frame(self, data: bytes):
        self.odom_encoder_left = self._parse_int32(data[:4])
        self.odom_encoder_right = self._parse_int32(data[4:8])
        self.get_motor_encoder_time =  self.node.get_clock().now().nanoseconds / 1e9
    
    def _handle_blade_frame(self, data: bytes):
        pass

    def _handle_state_request_frame(self, data: bytes):
        pass

    def _handle_unknown_command(self, cmd_id: int, data: bytes):
        print(f"[WARN] 收到未知命令，cmd_id: {cmd_id:#02x}, 长度: {len(data)}")

    @staticmethod
    def _parse_int32(data: bytes) -> int:
        # 大端解析有符号 int32
        return int.from_bytes(data, byteorder='big', signed=True)
    
    @property
    def serial_is_open(self):
        return self.stm32_serial.is_open

    @property
    def motor_encoder_counts(self):
        '''
        获取左右电机编码器值
        '''
        return (self.odom_encoder_left, self.odom_encoder_right)
    
    @property
    def motor_encoder_time(self):
        '''
        获取获取左右电机编码器值的时间戳
        '''
        return (self.get_uart_data_time,self.get_motor_encoder_time)
    
    #############
    #下行协议
    #############
    def write_to_serial(self,cmd_id: int, data: bytes):
        if not self.stm32_serial or not self.stm32_serial.is_open:
            self.node.get_logger().error("串口未打开，无法发送数据")
            return False
        with self.serial_lock:
            try:
                # 构建 payload（CmdID + Data）
                payload = bytes([cmd_id]) + data
                # self.node.get_logger().info(f"send payload: {' '.join(f'{b:02X}' for b in payload)}")
                length = len(payload) + 1
                # 计算 CRC（对 payload 进行异或）
                crc = self._calc_checksum(payload)
                # 构建完整帧：Header(2) + Length(1) + Payload(N) + CRC(1) + Tail(1)
                frame = self.FRAME_HEADER + bytes([length]) + payload + bytes([crc, self.FRAME_TAIL])
                # self.node.get_logger().info(f"send frame: {' '.join(f'{b:02X}' for b in frame)}")

                return self.stm32_serial.write(frame) == len(frame)
            except serial.SerialException as e:
                self.node.get_logger().error(f"Write failed: {e}")
                return False
    

    def send_encoder_data(self, left_encoder: int, right_encoder: int):
        encoder_data = left_encoder.to_bytes(2, 'big', signed=True) \
                       + right_encoder.to_bytes(2, 'big', signed=True)
        return self.write_to_serial(self.CMD_ENCODER, encoder_data)
          
        
       






    